package org.example.offical.doc;

import dev.langchain4j.data.message.*;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.memory.chat.TokenWindowChatMemory;
import dev.langchain4j.model.openai.OpenAiChatModelName;
import dev.langchain4j.model.openai.OpenAiTokenizer;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import dev.langchain4j.store.memory.chat.InMemoryChatMemoryStore;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author fish_temp_author
 * @since fish_temp_since
 */
public class ChatMemoryDemo {
    public static void main(String[] args) {
        SystemMessage systemMessage = SystemMessage.from("你是一位 experienced {{lang}} 开发工程师，请根据问题直接给出代码");
        UserMessage message1 = UserMessage.from(TextContent.from("你好，是1号"));
        UserMessage message2 = UserMessage.from(TextContent.from("你好，是2号"));
        UserMessage message3 = UserMessage.from(TextContent.from("你好，是3号"));
        UserMessage message4 = UserMessage.from(TextContent.from("你好，是4号"));
        UserMessage message5 = UserMessage.from(TextContent.from("你好，是5号"));
        SystemMessage systemMessage2 = SystemMessage.from("你现在是程序员鼓励师");


        List<ChatMessage> messageList = List.of(/*SystemMessage一旦添加，会一直存在于记忆中，但是只会有一条*/
                systemMessage,
                message1,
                // 后面的会覆盖前面的
                systemMessage2,
                message2, message3, message4, message5);

        MessageWindowChatMemory messageWindowChatMemory = MessageWindowChatMemory.builder()
                .maxMessages(4)
                .chatMemoryStore(new MyChatMemoryStore())
                .build();
        messageList.forEach(messageWindowChatMemory::add);
        System.out.println("messageWindowChatMemory ret: " + messageWindowChatMemory.messages());


        TokenWindowChatMemory tokenWindowChatMemory = TokenWindowChatMemory.builder()
                // 如果记忆中包含SystemMessage，且SystemMessage超过maxToken这个长度，会报错！
                .maxTokens(60, new OpenAiTokenizer(OpenAiChatModelName.GPT_4_O_MINI))
                // chatMemory不是chat history，如果是单实例部署的话感觉使用内存保存就可以
                .chatMemoryStore(new InMemoryChatMemoryStore())
                //.chatMemoryStore(new RedisChatMemoryStore("127.0.0.1", 6379, "", ""))
                .build();
        messageList.forEach(tokenWindowChatMemory::add);
        System.out.println("tokenWindowChatMemory ret: " + tokenWindowChatMemory.messages());
    }

    public static class MyChatMemoryStore implements ChatMemoryStore {

        Map<Object, List<ChatMessage>> memoryMap = new HashMap<>();

        /**
         * Implement getting all messages from the persistent store by memory ID.
         * ChatMessageDeserializer.messageFromJson(String)
         *
         * @param memoryId The ID of the chat memory.
         * @return
         */
        @Override
        public List<ChatMessage> getMessages(Object memoryId) {
            List<ChatMessage> chatMessages = memoryMap.getOrDefault(memoryId, List.of());
            // 使用占位符的方式打印
            System.out.println("getMessages - memoryId:, " + memoryId + " - getMessage :" + chatMessages);
            return chatMessages;
        }

        /**
         * 每次将消息添加到chatMemory都会调用这个方法，id用户区分用户会话
         * Implement updating all messages in the persistent store by memory ID.
         *
         * @param memoryId The ID of the chat memory.
         * @param messages List of messages for the specified chat memory, that represent the current state of the {@link ChatMemory}.
         *                 Can be serialized to JSON using {@link ChatMessageSerializer}.
         */
        @Override
        public void updateMessages(Object memoryId, List<ChatMessage> messages) {
            System.out.println("updateMessages - memoryId:" + memoryId + " - messages:" + messages);
            memoryMap.put(memoryId, messages);
        }

        /**
         * Implement deleting all messages in the persistent store by memory ID.
         * 每当调用ChatMemory.clear()时都会调用deleteMessages()方法
         *
         * @param memoryId The ID of the chat memory.
         */
        @Override
        public void deleteMessages(Object memoryId) {
            System.out.println("deleteMessages -  memoryId:" + memoryId);
            memoryMap.remove(memoryId);
        }
    }
}
